package org.example.visitor;

import com.neuronbit.parser.MySqlParser;
import org.antlr.v4.runtime.tree.ParseTree;

public class IdentifierVisitor extends com.neuronbit.parser.MySqlParserBaseVisitor {
    private TableInfo tableInfo = new TableInfo();

    @Override
    public Object visitTableName(com.neuronbit.parser.MySqlParser.TableNameContext ctx) {
        tableInfo.setTableName(ctx.getText());
        return super.visitTableName(ctx);
    }

    @Override
    public Object visitColumnName(com.neuronbit.parser.MySqlParser.ColumnNameContext ctx) {
        tableInfo.getColumnNames().add(ctx.getText());
        return super.visitColumnName(ctx);
    }

    @Override
    public Object visitSimpleIndexDeclaration(MySqlParser.SimpleIndexDeclarationContext ctx) {
        TableInfo.Index index = new TableInfo.Index();
        for (ParseTree child : ctx.children) {
            if (child instanceof com.neuronbit.parser.MySqlParser.UidContext) {
                index.setIndexName(child.getText());
            } else if (child instanceof com.neuronbit.parser.MySqlParser.IndexColumnNamesContext) {
                MySqlParser.IndexColumnNamesContext icnsContext = (MySqlParser.IndexColumnNamesContext)child;
                for (ParseTree childOfIcns : icnsContext.children) {
                    if (childOfIcns instanceof MySqlParser.IndexColumnNameContext) {
                        index.getIndexColumnNames().add(childOfIcns.getText());
                    }
                }
            }
        }
        tableInfo.getIndexes().add(index);
        return super.visitSimpleIndexDeclaration(ctx);
    }

    @Override
    public Object visitSpecialIndexDeclaration(MySqlParser.SpecialIndexDeclarationContext ctx) {
        TableInfo.Index index = new TableInfo.Index();
        for (ParseTree child : ctx.children) {
            if (child instanceof com.neuronbit.parser.MySqlParser.UidContext) {
                index.setIndexName(child.getText());
            } else if (child instanceof com.neuronbit.parser.MySqlParser.IndexColumnDefinitionContext) {
                for (ParseTree parseTree : ((MySqlParser.IndexColumnDefinitionContext) child).children) {
                    index.getIndexColumnNames().add(parseTree.getText());
                }
            }
        }
        tableInfo.getIndexes().add(index);
        return super.visitSpecialIndexDeclaration(ctx);
    }

    @Override
    public Object visitPartitionDefinitions(com.neuronbit.parser.MySqlParser.PartitionDefinitionsContext ctx) {
        for (ParseTree child :ctx.children) {
            if (child instanceof MySqlParser.PartitionComparisonContext) {
                for (ParseTree subChild : ((MySqlParser.PartitionComparisonContext) child).children) {
                    if (subChild instanceof MySqlParser.UidContext) {
                        tableInfo.getPartitions().add(subChild.getText());
                    }
                }
            }
        }
        return super.visitPartitionDefinitions(ctx);
    }

    public TableInfo getTableInfo() {
        return tableInfo;
    }
}
